# location based representation of object
import random
import uuid
from collections import defaultdict

import numpy as np
from skimage import measure

import utils as utils


class CauseEffectPosition:

    def __init__(self, causeEffectId, actionId, objInstId, positionDelta):
        self.id = causeEffectId
        self.actionId = actionId
        self.objInstId = objInstId
        self.positionDelta = positionDelta

    def __eq__(self, other):
        if self.actionId != other.actionId:
            return False
        # if self


class ObjectDescription:

    def __init__(self, objDescriptionId, coordinates, feature):
        self.id = objDescriptionId
        self.locations = set()
        self.feature = feature
        self.coordinates = coordinates
        self.locationToCoordinateMap = defaultdict(object)
        for coord in coordinates:
            loc = uuid.uuid1()
            self.locations.add(loc)
            self.locationToCoordinateMap[loc] = (coord[0:2])


class ObjectInstance:

    def __init__(self, objInstanceId, objDescriptionId, bbox):
        self.id = objInstanceId
        self.objDescriptionId = objDescriptionId
        self.bbox = bbox

    def getBboxTopLeft(self):
        return self.bbox[0:2]


class Agent:

    def __init__(self, worldDim):
        self.worldHeight = worldDim[0]
        self.worldWidth = worldDim[1]
        self.learnedObjectDescriptions = []  # need to be persisted
        self.matchedObjInst = []
        self.maxODID = 0
        self.maxObjInstId = 0
        self.maxCauseEffectId = 0
        self.currentScene = []
        self.previousScene = []
        self.actionSpace = [0, 1, 2, 3]
        self.previousAction = None
        self.learnedCauseEffects = []

    def createCurrentScene(self, observation):
        self.previousScene = self.currentScene
        self.currentScene = []
        encodedImage = utils.encodeImage(observation)
        label = measure.label(encodedImage, connectivity=2)
        props = measure.regionprops(label)
        # matching regions with learned objects
        for pp in props:
            firstCoord = pp.coords[0][0:2]
            feature = observation[firstCoord[0], firstCoord[1]]
            if len(self.learnedObjectDescriptions) == 0:
                # if there is no learned object description, just create a new one
                objDescrip = ObjectDescription(self.maxODID, pp.coords, feature)
                self.maxODID += 1
                self.learnedObjectDescriptions.append(objDescrip)
                # after create a new lod, create a new object instance
                objInst = ObjectInstance(self.maxObjInstId, objDescrip.id, pp.bbox)
                self.maxObjInstId += 1
                self.currentScene.append(objInst)
                continue
            foundMatchForTheRegion = False
            for objDescrip in self.learnedObjectDescriptions:
                # if match, create an object instance
                if self.matchDescription(pp.coords, feature, objDescrip):
                    objInst = ObjectInstance(self.maxODID, objDescrip.id, pp.bbox)
                    self.maxODID += 1
                    self.currentScene.append(objInst)
                    foundMatchForTheRegion = True
                    break
            if not foundMatchForTheRegion:
                # if do not match, create a new learned object description
                objDescrip = ObjectDescription(self.maxODID, pp.coords, feature)
                self.maxODID += 1
                self.learnedObjectDescriptions.append(objDescrip)
                objInst = ObjectInstance(self.maxObjInstId, objDescrip.id, pp.bbox)
                self.maxObjInstId += 1
                self.currentScene.append(objInst)

    def matchObjectBetweenScene(self):
        for objInstPrevious in self.previousScene:
            matchingObjInst = Agent.findNoChangesInTheScene(objInstPrevious, self.currentScene)
            if matchingObjInst:
                matchingObjInst.id = objInstPrevious.id
            else:  # not matching object that stay in the old place, check position change:
                positionChangedObjInst = Agent.findPositionChanges(objInstPrevious, self.currentScene)
                positionChangedObjInst.id = objInstPrevious.id

    def learnCauseEffect(self):
        if self.previousAction is None or len(self.previousScene) == 0:
            return
        for currenObjInst in self.currentScene:
            previousObjInst = self.findObjInstById(currenObjInst.id, self.previousScene)
            currentBbox = currenObjInst.bbox
            previousBbox = previousObjInst.bbox
            if currentBbox != previousBbox:
                existingCEP = self.getCauseEffectPosition(self.previousAction)
                if not existingCEP:
                    positionDelta = (np.array(currentBbox) - np.array(previousBbox))[0:2]
                    newcep = CauseEffectPosition(self.maxCauseEffectId, self.previousAction, currenObjInst.id, positionDelta)
                    self.maxCauseEffectId += 1
                    self.learnedCauseEffects.append(newcep)

    def getCauseEffectPosition(self, actionId):
        for cep in self.learnedCauseEffects:
            if cep.actionId == actionId:  # each atomic action has only one causal rule
                return cep
        return None

    @staticmethod
    def findObjInstById(objInstId, objInstList):
        for objInst in objInstList:
            if objInst.id == objInstId:
                return objInst

    @staticmethod
    def findNoChangesInTheScene(objInst, theScene):
        for objInstInTheScene in theScene:
            if objInst.objDescriptionId == objInstInTheScene.objDescriptionId and objInst.bbox == objInstInTheScene.bbox:
                return objInstInTheScene
        return None

    @staticmethod
    def findPositionChanges(objInst, theScene):
        for objInstInTheScene in theScene:
            if objInst.objDescriptionId == objInstInTheScene.objDescriptionId and objInst.bbox != objInstInTheScene.bbox:
                return objInstInTheScene
        return None

    def matchDescription(self, ppCoord, ppFeature, objectDescription):
        if utils.encodeRGB(objectDescription.feature) != utils.encodeRGB(ppFeature):
            return False
        if ppCoord.shape[0] != objectDescription.coordinates.shape[0]:
            return False
        displacement = ppCoord - objectDescription.coordinates
        firstDiff = utils.encodeCoord(displacement[0], self.worldWidth)
        for d in displacement:
            dEncoded = utils.encodeCoord(d, self.worldWidth)
            if firstDiff != dEncoded:
                return False
        return True

    def chooseRandomAction(self):
        action = random.randint(0, len(self.actionSpace)-1)
        # action = 0
        self.previousAction = action
        return action

    def getAction(self, targetObjInstId, targetBboxTopLeft):
        if len(self.learnedCauseEffects) < len(self.actionSpace):
            return self.chooseRandomAction()
        else:
            return self.plan(targetObjInstId, targetBboxTopLeft)

    def plan(self, targetObjInstId, targetBboxTopLeft):
        targetObjInst = self.findObjInstById(targetObjInstId, self.currentScene)
        currBboxTopLeft = targetObjInst.getBboxTopLeft()
        currManhattan = utils.manhattan(targetBboxTopLeft, currBboxTopLeft)
        availableAction = []
        # strategy one: try out actions in the mind
        for act in self.actionSpace:
            cep = self.findCauseEffectByActionId(act)
            updatedBboxTopLeft = currBboxTopLeft + cep.positionDelta
            updatedManhattan = utils.manhattan(targetBboxTopLeft, updatedBboxTopLeft)
            if updatedManhattan < currManhattan:
                availableAction.append(act)
        selectAct = random.randint(0, len(availableAction)-1)
        return availableAction[selectAct]

    def checkTargetFininsed(self, targetObjInstId, targetBboxTopLeft):
        targetObjInst = self.findObjInstById(targetObjInstId, self.currentScene)
        currBboxTopLeft = targetObjInst.getBboxTopLeft()
        currManhattan = utils.manhattan(targetBboxTopLeft, currBboxTopLeft)
        if currManhattan == 0:
            print('reached target, no action')
            return True
        else:
            return False

    def findCauseEffectByActionId(self, actionId):
        for cep in self.learnedCauseEffects:
            if cep.actionId == actionId:
                return cep

    @staticmethod
    def printScene(theScene):
        print('there are', len(theScene), 'objects in the scene')
        for ith, o in enumerate(theScene):
            print(str(ith) + ' object instance Id is ', o.id, '; object description Id is', o.objDescriptionId)
